#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy, sys
import moveit_commander
from moveit_msgs.msg import RobotTrajectory
from trajectory_msgs.msg import JointTrajectoryPoint
from geometry_msgs.msg import PoseStamped, Pose, PointStamped
from vision_msgs.msg import Detection2D
import tf2_ros
import tf2_geometry_msgs

class MoveItGrasp:
    def __init__(self):
        # Initialize ROS node
        rospy.init_node('moveit_grasp_demo')
        r = rospy.Rate(10)
        rospy.on_shutdown(self.shutdown)
                
        # Initialize MoveIt commander
        moveit_commander.roscpp_initialize(sys.argv) 
        self.arm = moveit_commander.MoveGroupCommander('manipulator')        
        self.end_effector_link = self.arm.get_end_effector_link()
        self.reference_frame = 'world'
        self.arm.set_pose_reference_frame(self.reference_frame)

        # Initialize TF2
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

        # MoveIt planning parameters
        self.arm.set_planning_time(10.0)
        self.arm.set_num_planning_attempts(50)
        self.arm.set_max_velocity_scaling_factor(0.3)
        self.arm.set_max_acceleration_scaling_factor(0.3)
        self.arm.allow_replanning(True)
        self.arm.set_goal_position_tolerance(0.01)
        self.arm.set_goal_orientation_tolerance(0.05)
        self.arm.set_planner_id("RRTConnect")  # Specify planner
        rospy.loginfo("Planner ID set to: %s", self.arm.get_planner_id())

        # Grasp control variables
        self.cnt = 0
        self.stable_cnt = 10
        self.grasp_status = False
        self.grasp_obj = None

        # Gripper offset (from end-effector to grasp point, in meters)
        self.gripper_offset_x = rospy.get_param('~gripper_offset_x', 0.0)
        self.gripper_offset_y = rospy.get_param('~gripper_offset_y', 0.0)
        self.gripper_offset_z = rospy.get_param('~gripper_offset_z', 0.0)

        # Subscriber for object coordinates
        self.object_sub = rospy.Subscriber("/object_point_info", Detection2D, self.ObjectCallback)

        # Move to initial pose
        self.init_pose()
        rospy.sleep(2)

        while not rospy.is_shutdown():
            if self.grasp_status and self.grasp_obj is not None:
                rospy.loginfo("start grasp object")
                x, y, z = self.camera_to_gripper_pose(
                    self.grasp_obj.x, self.grasp_obj.y, self.grasp_obj.z)
                if any(v is None for v in (x, y, z)):
                    rospy.logerr("Invalid grasp coordinates: x=%s, y=%s, z=%s", x, y, z)
                    self.grasp_status = False
                    self.cnt = 0
                    continue
                rospy.loginfo("move to grasp: %f, %f, %f", x, y, z)
                self.move_pose(x, y, z) 
                rospy.sleep(1)
                self.init_pose()
                rospy.sleep(2)
                self.cnt = 0
                self.grasp_status = False
            r.sleep()

    def ObjectCallback(self, msg):
        obj_info = msg
        self.cnt += 1
        if self.cnt > self.stable_cnt and not self.grasp_status:
            self.grasp_status = True
            self.grasp_obj = obj_info.results[-1].pose.pose.position
            rospy.loginfo("Received target point: x=%.3f, y=%.3f, z=%.3f", 
                          self.grasp_obj.x, self.grasp_obj.y, self.grasp_obj.z)

    def camera_to_gripper_pose(self, x, y, z):
        """
        Transform object coordinates (mm, camera frame) to gripper pose (m, base_link frame).
        """
        # Create PointStamped in camera frame (convert mm to m)
        point_camera = PointStamped()
        point_camera.header.frame_id = "camera_link"
        point_camera.header.stamp = rospy.Time(0)
        point_camera.point.x = x / 1000.0
        point_camera.point.y = y / 1000.0
        point_camera.point.z = z / 1000.0

        try:
            # Transform to base_link frame
            point_base = self.tf_buffer.transform(point_camera, "world", 
                                                timeout=rospy.Duration(1.0))
            # Apply gripper offset
            grasp_x = point_base.point.x + self.gripper_offset_x
            grasp_y = point_base.point.y + self.gripper_offset_y
            grasp_z = point_base.point.z + self.gripper_offset_z
            rospy.loginfo("Transformed target point: x=%.3f, y=%.3f, z=%.3f", 
                          grasp_x, grasp_y, grasp_z)
            return grasp_x, grasp_y, grasp_z
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, 
                tf2_ros.ExtrapolationException) as e:
            rospy.logerr("TF transform failed: %s", str(e))
            return None, None, None

    def move_pose(self, x, y, z):
        target_pose = PoseStamped()
        target_pose.header.frame_id = self.reference_frame
        target_pose.header.stamp = rospy.Time.now()     
        current_pose = self.arm.get_current_pose(self.end_effector_link)
        target_pose.pose.position.x = x
        target_pose.pose.position.y = y
        target_pose.pose.position.z = z   
        target_pose.pose.orientation = current_pose.pose.orientation
        self.arm.set_start_state_to_current_state()
        self.arm.set_pose_target(target_pose, self.end_effector_link)
        
        # Plan the motion
        traj = self.arm.plan()
        # Check if planning succeeded (trajectory is not None and has points)
        if traj and (hasattr(traj, 'joint_trajectory') and traj.joint_trajectory.points):
            self.arm.execute(traj, wait=True)
            rospy.loginfo("Successfully executed grasp motion")
        else:
            rospy.logwarn("Planning failed: No valid trajectory found")

    def init_pose(self):
        self.arm.set_named_target("look")
        success = self.arm.go(wait=True)
        if success:
            rospy.loginfo("Moved to look state successfully")
        else:
            rospy.logwarn("Failed to move to look state")

    def shutdown(self):
        moveit_commander.roscpp_shutdown()
        moveit_commander.os._exit(0)
        rospy.logwarn("now will shutdown grasp_node ...")

if __name__ == "__main__":
    MoveItGrasp()
